package com.idp.metadata;

import org.opensaml.Configuration;
import org.opensaml.saml2.metadata.*;
import org.opensaml.xml.XMLObject;
import org.opensaml.xml.io.Unmarshaller;
import org.opensaml.xml.io.UnmarshallerFactory;
import org.opensaml.xml.io.UnmarshallingException;
import org.opensaml.xml.security.credential.UsageType;
import org.opensaml.xml.security.x509.BasicX509Credential;
import org.opensaml.xml.signature.X509Certificate;
import org.opensaml.xml.util.Base64;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.xml.sax.SAXException;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.util.*;

public class MetadataBean {
    private static final Logger logger = LoggerFactory.getLogger(MetadataBean.class);
    private String entityId;
    private SSODescriptor ssoDescriptor;
    private List<BasicX509Credential> signingCredList;
    private List<BasicX509Credential> encryptionCredList;
    private List<NameIDFormat> nameIDFormatList;
    private Map<String, List<Endpoint>> endpointMap;

    public String getEntityId() {
        return entityId;
    }

    public List<BasicX509Credential> getSigningCredList() {
        return signingCredList;
    }

    public MetadataBean(String filePath) {
        initMetadata(filePath);
    }

    public void initMetadata(String filePath) {
        String metadataContent = getMetadataString(filePath);

        if (StringUtils.isEmpty(metadataContent)) {
            System.out.println("error : metadataContent is empty");
            return;
        }
        EntityDescriptor entityDescriptor = loadEntityDescriptorFromString(metadataContent);

        if (null != entityDescriptor.getIDPSSODescriptor(MetadataConstants.SAML2_SUPPORTED_PROTOCOL)) {
            this.ssoDescriptor = entityDescriptor.getIDPSSODescriptor(MetadataConstants.SAML2_SUPPORTED_PROTOCOL);

            this.entityId = entityDescriptor.getEntityID();
            try {
                this.signingCredList = generateIDPCredList();
            } catch (Exception e) {
                logger.error("", e);
            }
            try {
                this.encryptionCredList = generateIDPCredList();
            } catch (Exception e) {
                logger.error("", e);
            }
        } else {
            this.ssoDescriptor = entityDescriptor.getSPSSODescriptor(MetadataConstants.SAML2_SUPPORTED_PROTOCOL);
            this.entityId = entityDescriptor.getEntityID();
            this.signingCredList = generateSPCredList(UsageType.SIGNING);
            this.encryptionCredList = generateSPCredList(UsageType.ENCRYPTION);
            this.nameIDFormatList = ssoDescriptor.getNameIDFormats();
            SPSSODescriptor spssoDescriptor = (SPSSODescriptor) this.ssoDescriptor;
            this.endpointMap = new HashMap<>();

            this.endpointMap.put(AssertionConsumerService.DEFAULT_ELEMENT_LOCAL_NAME,
                    covertToEndpoints(spssoDescriptor.getAssertionConsumerServices()));

            this.endpointMap.put(ArtifactResolutionService.DEFAULT_ELEMENT_LOCAL_NAME,
                    covertToEndpoints(spssoDescriptor.getArtifactResolutionServices()));

            this.endpointMap.put(SingleLogoutService.DEFAULT_ELEMENT_LOCAL_NAME,
                    covertToEndpoints(spssoDescriptor.getSingleLogoutServices()));

            this.endpointMap.put(ManageNameIDService.DEFAULT_ELEMENT_LOCAL_NAME,
                    covertToEndpoints(spssoDescriptor.getManageNameIDServices()));
        }
    }

    private <T extends Endpoint> List<Endpoint> covertToEndpoints(List<T> services) {
        return new ArrayList<>(services);
    }

    private List<BasicX509Credential> generateSPCredList(UsageType usageType) {
        List<KeyDescriptor> keyDescriptors = ssoDescriptor.getKeyDescriptors();
        if (null == keyDescriptors) {
            throw new RuntimeException("sp keyDescriptors is null");
        }
        BasicX509Credential basicX509Credential = new BasicX509Credential();
        for (KeyDescriptor keyDescriptor : keyDescriptors) {
            if (keyDescriptor.getUse() == usageType) {
                X509Certificate x509Certificate = keyDescriptor.getKeyInfo().getX509Datas().get(0).getX509Certificates()
                        .get(0);
                java.security.cert.X509Certificate certificate = covertSamlCertToJavaCert(x509Certificate);
                basicX509Credential.setEntityCertificate(certificate);
                break;
            }
        }
        List<BasicX509Credential> basicX509Credentials = new ArrayList<BasicX509Credential>();
        basicX509Credentials.add(basicX509Credential);
        return basicX509Credentials;
    }

    private java.security.cert.X509Certificate covertSamlCertToJavaCert(X509Certificate x509Certificate) {
        CertificateFactory certificateFactory = null;
        try {
            certificateFactory = CertificateFactory.getInstance("X.509");
        } catch (CertificateException e) {
            logger.error("", e);
        }
        String certificateString = x509Certificate.getValue();
        byte[] bytes = Base64.decode(certificateString);
        ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(bytes);
        java.security.cert.X509Certificate certificate = null;
        try {
            certificate = (java.security.cert.X509Certificate) Objects.requireNonNull(certificateFactory).generateCertificate(byteArrayInputStream);
        } catch (CertificateException e) {
            logger.error("", e);
        }
        return certificate;
    }

    private static List<BasicX509Credential> generateIDPCredList() {
        String fileName = IDPMetadata.KEY_STORE_PATH;
        char[] password = IDPMetadata.KEY_STORE_PASS.toCharArray();
        java.security.cert.X509Certificate certificate = null;
        PrivateKey privateKey = null;
        try (
                InputStream inputStream = new FileInputStream(fileName)
        ) {
            KeyStore keystore = KeyStore.getInstance(KeyStore.getDefaultType());
            keystore.load(inputStream, password);
            String aliasName = IDPMetadata.ALIAS_NAME;
            certificate = (java.security.cert.X509Certificate) keystore.getCertificate(aliasName);
            privateKey = (PrivateKey) keystore.getKey(aliasName, password);

        } catch (Exception e) {
            logger.error("", e);
        }
        List<BasicX509Credential> basicX509Credentials = new ArrayList<BasicX509Credential>();
        BasicX509Credential basicX509Credential = new BasicX509Credential();
        basicX509Credential.setEntityCertificate(certificate);
        basicX509Credential.setPrivateKey(privateKey);
        basicX509Credentials.add(basicX509Credential);

        return basicX509Credentials;
    }


    private static String getMetadataString(String filePath) {
        String str = "";
        File spMetadataFile = new File(filePath);
        try (
                FileInputStream in = new FileInputStream(spMetadataFile)
        ) {
            int size = in.available();
            byte[] buffer = new byte[size];
            int read = in.read(buffer);
            str = new String(buffer, StandardCharsets.UTF_8);
        } catch (IOException e) {
            logger.error("", e);
        }
        return str;
    }

    private EntityDescriptor loadEntityDescriptorFromString(String content) {
        ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(content.getBytes());
        return (EntityDescriptor) unmarshallXMLObject(byteArrayInputStream);
    }

    private static XMLObject unmarshallXMLObject(InputStream inputStream) {
        try {
            DocumentBuilderFactory documentBuilderFactory = DocumentBuilderFactory.newInstance();
            documentBuilderFactory.setFeature("http://apache.org/xml/features/disallow-doctype-decl", true);
            documentBuilderFactory.setNamespaceAware(true);
            DocumentBuilder documentBuilder = documentBuilderFactory.newDocumentBuilder();
            Document document = documentBuilder.parse(inputStream);
            Element element = document.getDocumentElement();
            UnmarshallerFactory unmarshallerFactory = Configuration.getUnmarshallerFactory();
            Unmarshaller unmarshaller = unmarshallerFactory.getUnmarshaller(element);
            return unmarshaller.unmarshall(element);
        } catch (ParserConfigurationException | SAXException | IOException | UnmarshallingException e) {
            logger.error("", e);
        } finally {
            if (inputStream != null) {
                try {
                    inputStream.close();
                } catch (IOException e) {
                    logger.error("", e);
                }
            }
        }
        return null;
    }
}
